package org.example.myleet.p1337;

import java.util.Arrays;
import java.util.Comparator;

public class Solution {
    /**
     * 利用二分查找加速军人数求和
     */
    public int[] kWeakestRows(int[][] mat, int k) {
        for (int i = 0; i < mat.length; ++i) {
            mat[i][0] = getSoldierSum(mat[i]);
            mat[i][1] = i;
        }
        Arrays.sort(mat, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                if (o1[0] != o2[0]) {
                    return o1[0] - o2[0];
                }
                return o1[1] - o2[1];
            }
        });
        int[] result = new int[k];
        for (int i = 0; i < k; ++i) {
            result[i] = mat[i][1];
        }
        return result;
    }

    /**
     * 题目要求军人全部排在前面
     */
    private int getSoldierSum(int[] arr) {
        int l = 0, r = arr.length - 1;
        while (r - l > 1) {
            int m = (r + l) / 2;
            if (arr[m] == 1) {
                l = m;
            } else {
                r = m;
            }
        }
        if (arr[l] == 0) {
            return 0;
        }
        if (arr[r] == 1) {
            return r + 1;
        }
        return l + 1;
    }
}
